Per-catchment CFE DA: true EnKF (N=20) + Vrugt R + mass-conserving cascade - #7

Test loop (run_testing_period) now runs N=20 parallel CFE members with
multiplicatively perturbed precipitation (15%) and PET (10%) per member,
plus 5% initial state perturbation. At each hour the ensemble of forecast
discharges produces Pyy and the state-to-Q cross-covariances; one Kalman
gain per state is derived from the ensemble (Burgers/Evensen 1998), the
observation is perturbed N times, and every member is updated
independently. The output time series is the ensemble mean. This replaces
the 30/15/20/35 hardcoded split and the (0.3 * Q_sim)^2 forecast-variance
heuristic with data-driven covariance.

Calibration loop keeps the single-trajectory heuristic update via the
renamed update_states_single() helper, so DDS cost is unchanged.

Mass-conserving cascade preserved per member:
  positive : soil -> Nash[0], GW -> Nash[1]
  negative : Nash[1] -> Nash[0] -> soil -> GW -> loss (logged in mm)

New EnKFAssimilator constructor args: precip_perturb_frac, pet_perturb_frac,
init_state_perturb_frac, rng_seed. Diagnostic line at end of test run shows
N, mean innovation, average Pyy, per-state mean Kalman gains, and total
mass lost at GW underflow.

Old docstring described an EnKF that was not actually implemented (no
ensemble), described a single state-update path, and used the wrong
script name in the usage example. The new text accurately states:

  - the calibration loop uses a single-trajectory heuristic
  - the test loop uses a true stochastic EnKF over N ensemble members
  - 4 CFE states are updated with a mass-conserving cascade
  - all three obs CSV column layouts are auto-handled
  - the production usage example points at the right paths

Removed boilerplate "literature improvements" numbers in favor of the
measured result on this basin (mean Test KGE 0.692 -> 0.823, 21
catchments, Helene year).

The recorded Q_sim at time t is the pre-analysis forecast (one-step-ahead
under the most recent analysis). DA's effect on the series propagates
forward through carried states, not by overwriting Q_sim(t). This is the
standard sequential-filter verification setup; adding a comment so future
readers don't mistake the order of operations for a correctness issue.

Diagnostic on the previous run (sigma=0.005 across all states) showed that:
  - Nash gains jumped from ~1e-87 to ~1e-4 (fix worked)
  - Soil/GW gains also went up significantly
  - BUT Test KGE dropped 0.882 -> 0.835 and mass lost at GW underflow
    jumped 8.5x (9.9 -> 84.2 mm)

Root cause: 0.5%/hr is too aggressive for slow-evolving soil/GW. Over
9000 hours, soil values do a random walk with cumulative std ~47% of
nominal — way more than physics warrants. GW underflow becomes frequent.

Fix: per-state sigma calibrated to physical response time, matching
NWC-CUAHSI/data_assimilation_with_bmi reference implementation:
  soil  : 0.002 (0.2%)   — slow-evolving, small noise sufficient
  GW    : 0.0015 (0.15%) — slowest state
  Nash  : 0.005 (0.5%)   — needs more; spread doesn't develop naturally

Nash still gets the additive floor so noise survives at zero.

Old single knob 'process_noise_frac' removed; replaced by three per-state
constructor args. Set any to 0 to disable that state's process noise.

Per Dr Frame's suggestion on the Slack thread: use Vrugt et al. 2005 (SODA)
heteroscedastic observation error formulation, scaled by the kriging variance,
so that DA can fire at low flow without being dominated by the large raw
kriging variance values in Kunal's obs files.

  R(t) = (alpha * y_obs(t))^2 + scale * sigma^2_krig(t)

Two terms:
  - (alpha * y_obs)^2  : Vrugt flow-magnitude scaling, alpha=0.10 default
  - scale * sigma^2    : per-hour kriging contribution, scale=0.001 default

Numbers for this basin at typical y_obs ~ 0.2 mm/h, sigma^2_krig ~ 4:
  (0.1 * 0.2)^2 + 0.001 * 4 = 4e-4 + 4e-3 = 0.0044
vs raw kriging variance = 4 (which collapses K to ~0).

New CLI flags:
  --use-vrugt-r    : enable Vrugt R (default off)
  --vrugt-alpha    : relative-error fraction (default 0.10)
  --vrugt-scale    : kriging-variance scaling (default 0.001)

Pre-computed at obs-load time and stored in self.obs_var_dict, so both
DA paths (test loop true EnKF and cal loop heuristic) automatically use
the new R when --use-vrugt-r is enabled. Default OFF preserves prior
runs' behavior exactly.

After comparing four configurations (no PN/single R, per-state PN/raw R,
per-state PN/Vrugt R) across all 21 catchments at gauge 03463300, the
Vrugt R configuration is shipped as the production default:

  Mean Test KGE  : 0.817 (Vrugt) vs 0.823 (raw R)
  Best-worst spread : 0.175 (Vrugt) vs 0.220 (raw R) - much tighter
  cat-1016306 (the persistent regressor) : 0.66 -> 0.74 - recovered

Mean is 0.006 lower but the distribution is meaningfully more uniform,
and the worst catchment finally beats Run 3 baseline. The Vrugt
formulation is also the more defensible scientific story (heteroscedastic
R, kriging variance actually used, matches Dr Frame's Slack suggestion).

Changes:
  - Class default: use_vrugt_r=True
  - CLI flag flipped: --use-vrugt-r -> --no-vrugt-r (default ON; flag to disable)
  - Top docstring + USAGE example updated; production obs dir is now
    catchment_ts_03463300_with_variance so kriging variance enters R(t)
  - All fallback defaults inside EnKFAssimilator() construction sites
    set to True.

plot_da_v2_vrugt_helene_grid.py — generates a 3x7 grid comparing Run 3
(no DA) vs DA v2 (true EnKF + Vrugt R, N=20) vs Qkrig observations for
all 21 sub-catchments during the Hurricane Helene week (Sep 20 - Oct 5,
2024). Per-subplot KGE is displayed in the legend. Same template as the
prior DA v3/v4 grid plots.

da_v2_vrugt_kge_comparison.md — per-catchment KGE table for the
production DA v2 Vrugt R configuration vs Run 3 baseline. Mean Test KGE
improves from 0.692 to 0.817 (+0.125). 20/21 catchments improve;
cat-1016306 is the lone regressor (-0.033). Spread tightens from
0.325 to 0.175.

Replaces da_v2_vrugt_kge_comparison.md (which was markdown) with an
actual rendered table image.

helene_da_v2_vrugt_vs_run3_grid.png — 3x7 grid of all 21 sub-catchments
during Hurricane Helene week, comparing Run 3 baseline (thick blue),
DA v2 production config (thin red), and Qkrig observations (black).
Per-subplot KGE shown in the legend. Generated by
plot_da_v2_vrugt_helene_grid.py.

da_v2_vrugt_kge_table.png — color-coded per-catchment KGE comparison
table. Green = improvement, yellow = neutral, red = regression. Mean
row highlighted. 20/21 catchments improve; cat-1016306 is the lone
regressor (-0.033). Mean Test KGE goes from 0.692 (Run 3) to 0.817
(DA v2 Vrugt R).

Same 3x7 grid as plot_da_v2_vrugt_helene_grid.py (Run 3 vs DA v2 Vrugt R
vs Qkrig obs), but with a tighter 4-day window around the Helene peak so
the timing and amplitude detail at the rising/falling limbs is visible.
Daily x-ticks plus 6-hour minor ticks.

3x7 grid of all 21 sub-catchments showing the 4-day Helene peak window.
Same Run 3 vs DA v2 Vrugt R vs Qkrig comparison as the 16-day grid, but
tightened to make peak timing and amplitude visible.

Generated by plot_da_v2_vrugt_helene_zoomed_grid.py.

plot_perturbation_sensitivity_spread.py — cleaner alternative to the
spaghetti plot. For each catchment and each of the three sensitivity
sub-experiments, computes the hourly std-dev across the 20 members and
plots all three curves on the same axes. Directly answers 'which source
contributes most to spread' without the visual clutter of overlapping
member bundles.

run_production_per_member.py — production-equivalent run for ONE
catchment (init + forcing + process + DA all active, Vrugt R) but saves
per-member streamflow trajectories instead of only the ensemble mean.
Imports EnKFAssimilator from calibrate_catchment_cfe_da_v2.py to
guarantee identical DA math. Skips spinup to match the sensitivity-run
setup, so per-member comparisons against the sensitivity CSVs are
apples-to-apples.

plot_per_member_factor_decomposition.py — 20-panel figure for one
catchment. Each panel shows ONE ensemble member's actual production
forecast (purple thick) decomposed against its trajectory under each
isolated perturbation source (red = init only, blue = forcing only,
green = process noise only), plus the Qkrig observation. Helene 4-day
peak window. First-cut visualization (member identity is column-name
only across the four CSVs; not matched-seed).

Plots all 20 production-per-member streamflow trajectories overlaid on
one time-series panel, styled after the Battula et al. ensemble-forecast
figure (50/100/200 km variogram-range panels):
  - Each member as a thin turbo-colored line, labeled with index + peak
  - Hurricane Helene window (Sep 26 12:00 - Sep 28 00:00) highlighted as
    a pink shaded vertical band
  - Qkrig observation overlaid as a thick black line on top
  - Plot window Sep 20 - Oct 5 so the storm sits in context
  - Two outputs: linear-y and log-y versions

Consumes the production-per-member CSV produced by
run_production_per_member.py. No GPU re-run required if that CSV already
exists for the target catchment.

…stic plot

run_production_per_member.py — now saves three additional outputs per run:
  - <cat>_production_per_member_precip.csv (perturbed precip per member)
  - <cat>_production_per_member_pet.csv    (perturbed PET per member)
  - <cat>_production_per_member_initial_states.json (4 states per member at t=0)

plot_per_member_input_output_diagnostic.py — new plot script.
Layout: 20 rows x 4 columns. Each row is one ensemble member.
  col 1: perturbed precipitation (bars)
  col 2: perturbed PET (line)
  col 3: simulated streamflow + Qkrig obs
  col 4: initial-state values as annotation
Helene 4-day peak window (Sep 24-28, 2024).

Answers 'why did THIS member produce THIS forecast' by tracing inputs
to output, rather than running counterfactual sub-experiments.

Each member is now ONE combined panel instead of 4 separate columns.
Internal layout per panel:
  top    : perturbed precip (blue bars, left axis) + perturbed PET
           (orange line on twinx, separate scale)
  bottom : simulated Q (purple) vs Qkrig obs (black dashed), shared x-axis
  title  : initial-state values for that member (soil, GW, Nash[0], Nash[1])

Grid: 5 rows x 4 cols = 20 panels. All panels share the same y-axis ranges
across precip, PET, and Q for direct visual comparison between members.

Each input/output stream is on its own y-axis with a label and matching
color (P blue, PET orange, Q purple) so even at thumbnail size you can
read which trace is which.

Two layout tweaks to the diagnostic plot:
  - Split each panel title into two lines (member name on top, initial-state
    values below) so they no longer collide with the adjacent panel's title.
  - Bump hspace inside each member cell (precip+PET above, Q below) from
    0.05 to 0.20 so the bottom Q-axis ticks no longer brush the top of the
    next subpanel down.
  - Slightly increase outer GridSpec hspace to accommodate the taller
    2-line titles.

For one catchment, plot three category-organized ensemble groups:
  1. Initial states only       (red)
  2. Meteorological forcings   (blue)
  3. Hydrological states only  (green)

Each category is rendered as a 10th-90th percentile shaded band across
its 20 members, plus a thicker median line in the same color. Qkrig
observation overlaid in black. Hurricane Helene peak window shaded in
pink. Paper-style figure inspired by Battula et al. ensemble-forecast
plot.

Consumes the existing per-source sensitivity CSVs - no new GPU run
needed. Outputs both linear-y and log-y versions of the figure.

…nvelope

Production-magnitude perturbations produce narrow ensemble spread at the
Helene peak. The 10-90 percentile bands were too tight to be visually
distinguishable. Min-max envelope across the 20 members is the widest
possible band from the existing data and reveals the spread structure
more clearly, especially for the forcing-only category which has the
biggest dispersion at the peak.

Sensitive to outliers, but that is part of the story: a few members
hitting bigger perturbation factors push the envelope up at the peak.

Restructures da_methods/ into subdirectories matching the Qkrig DA repo
design: 1_calibrate, 2_assimilation, 3_routing, 4_evaluation.

Adds new scripts for the full Hurricane Helene probabilistic forecasting
pipeline: run_perturbation_da_on.py (2a/2b DA arms), run_crossed_ensemble.py
(600-member 30x20 design), run_route_crossed_ensemble.py (t-route routing
of ensemble to gauge), batch runners for all 21 catchments, and evaluation
plots for Vrugt vs no-Vrugt comparison and routed ensemble vs USGS gauge.

Also adds da_methods/README.md documenting the end-to-end data flow and
run order from calibration through gauge evaluation.

Copies run_perturbation_sensitivity.py (DA-off, 20-member ensemble, 3 sources:
init/forcing/process) from F1 into F2 and F5. Adds matching batch scripts
batch_run_sensitivity_f2.sh and batch_run_sensitivity_f5.sh pointing at each
folder's OUT_DIR and correct OBS_DIR (gapfilled for F2, rekrig for F5).

… F4/F5 4b scripts

route_ensemble_f1/f2/f4.sh route all 600 crossed-ensemble members through
Muskingum-Cunge to produce routed_crossed_ensemble.parquet (same pattern as
existing route_ensemble_f5.sh). Fixes wrong _routed suffix in F1_DIR path
in run_4b_f4.sh and run_4b_f5.sh.

F4: add run_perturbation_sensitivity.py and batch_run_sensitivity_f4.sh.
F5: add new_EnKF.py, input_enkf_new.json, run_crossed_ensemble.py,
run_perturbation_da_on.py (copied from F4 — same --direct-variance logic,
different OBS_DIR pointing at rekrig obs), batch_run_crossed_f5.sh,
batch_run_leadtime_f5.sh, batch_run_da_on_all_cats.sh, route_leadtime_f5.sh.
All four folders now have complete spec coverage for items 2 and 3.

… forcing arm

F1: helene_sensitivity_spaghetti/spread main+appendix (from run_4b_f1.sh).
F2: helene_ensemble_twopanel, helene_ensemble_vs_usgs, routed_ensemble_vs_usgs
    full+peak (from run_4b_f2.sh), cat-1016300_2a_forcing_arm_helene.
F4: cat-1016300_2a_forcing_arm_helene.
F5: cat-1016300_2a_forcing_arm_helene.
All four folders now have 2a+2b arm, 4b ensemble, and spec item 4b figures.

…ython scripts

All 4 experiment folders (F1 vrugt, F2 fixed-R, F3 dyn-vrugt-seeded, F4 dyn-variance):
- 3_routing/: route_det, route_leadtime, route_ensemble scripts (12 new scripts)
- 4a/: run_4a_f*.sh for all 4 folders
- 4b/: run_4b_f*.sh, run_4b_crossed_f*.sh, plot_routed_ensemble_vs_usgs.py for all 4 folders
- 4c/: run_4c_f1.sh, run_4c_f2.sh (F3/F4 already had these)
- 2_assimilation/: run_lead_time_forecast_sweep.py (F1), run_perturbation_sensitivity.py
  and run_production_per_member.py (F2/F3/F4)

- batch_run_f5_20pct.sh: production DA run (EnKF N=20, --no-vrugt-r)
  using re-kriged variance obs and held-out calibration params
- batch_run_f5_analysis_20pct.sh: perturbation arms (Stage A) and
  lead-time forecast sweep (Stage B) for uncertainty quantification
- route_det_f5_20pct.sh: deterministic T-route routing to gauge 03463300
- batch_run_openloop_20pct.sh: open-loop baseline (no DA) for 20% holdout
- route_openloop_20pct.sh: T-route routing for open-loop baseline
- compare_all_folders.py: updated to include F5 in metrics table and plots

Routed results: Full KGE=+0.317, NSE=+0.556; Helene KGE=+0.286, NSE=+0.461

Adds arm plots (2a, 2b, 2ab), lead time decay, open-loop comparison,
and multi-folder view plots for the F5 re-kriged σ² formulation.
Fixes F2 routing script to write output to a separate directory.

…holdout)

- update view1_full_log_f2_vs_ol.png: title now reads
  'F5 (best NSE) and F2 (best KGE) vs Open loop'
- add folder5_rekrig_variance_direct/4_evaluation/4a_error_decay/:
    plot_lead_time_nse_gauge_f5.py  — NSE vs lead time (pooled + regime split)
    run_4a_f5.sh                    — shell driver for the server

- lead_time_nse_gauge_f5_pooled.png: pooled NSE vs lead hour (DA vs open loop)
- lead_time_nse_gauge_f5_by_regime.png: NSE by regime (Helene / storm / low-flow)
- fix int(np.where(...)[0]) → int(np.where(...)[0][0]) to silence DeprecationWarning

Per-init-time coloured fan (min-max shading + median) for Sep 24-30 2024,
grand mean, open-loop grand median, and USGS obs at routed focal outlet.

…end x-axis

…ay in fan plot

…AD_HOURS=120)

Forecast window is 18 hours (not 120); reverts the variable to
match the intended short-range evaluation horizon for this holdout.

With FORECAST_LEAD_HOURS=18, ticks every 1h (0..18) give a
readable x-axis; the previous 12h spacing left only 0 and 12.

Plots DA ensemble mean vs Qkrig obs at catchment level (mm/h)
before routing, so DA skill can be assessed independently of
T-Route. Full period + Helene zoom panels, NSE in title.

Plots T-Route output (m³/s) vs USGS obs at gauge 03463300.
Full period + Helene zoom, NSE and KGE in title.

Figure was generated using wrong obs source — needs to be
regenerated with 20% holdout kriging obs directory.

Generated using 20% holdout kriging obs directory
(catchment_ts_03463300_dynamic_variance_rekrig), replacing
the previously removed incorrect version.

Routes 21-catchment hydro-state arm CSVs (mm/h) through T-Route
Muskingum-Cunge to gauge 03463300 (m3/s) for each Helene init day,
then overlays open-loop grand median from parquet and USGS obs.
Saves cat-1016300_2b_hydro_arm_helene.png with open loop dashed line.

… 21 catchments)